from efficientnet_pytorch import EfficientNet


def create_efficientnet(dropout_rate=0.1):
    # 调整输入分辨率适配64x64
    model = EfficientNet.from_name('efficientnet-b0',
                                   in_channels=3,
                                   num_classes=200,
                                   image_size=64,
                                   dropout_rate=dropout_rate
                                   )
    for block in model._blocks:
        if hasattr(block, '_dropout'):
            block._dropout.p = dropout_rate

    return model

if __name__ == '__main__':
    model = create_efficientnet(dropout_rate=0.1)
    print(model)